# -*- coding: utf-8 -*-
# Including simulation data & experimental data
import scipy.io as scio
import numpy as np
import pandas as pd
import os
import fnmatch
import pymongo
from typing import Generator

def is_sublist(sublist:list, list_par:list):
    bool_list = []
    for ele in sublist:
        bool_list.append(ele in list_par)
    if np.all(bool_list):
        return True
    else:
        return False

def readDirectory(path:str, col_sour:pymongo.collection.Collection):
    period = 0.02
    sample_step = 0.001
    # 0.2s 时出现故障
    delimiter = 0.2
    if path[-2:].isdigit():
        labels = [path[-1], path[-2]]
    elif not path[-1].isdigit():
        labels = ["0"]
    else:
        labels = [path[-1]]
    # set data target label
    dataDict = {}
    for file in os.listdir(path):
        if fnmatch.fnmatch(file, "*.mat"):
            matDict = scio.loadmat(os.path.join(path, file))
            for k, v in matDict.items():
                if type(v) == np.ndarray:
                    time = v[:,0]
                    data = v[:,1]
                    maxTime = 0.2 + period
                    # 0.04 is two period of sample period.
                    minTime = maxTime - 3 * period
                    # 考虑尾端问题
                    sample_points_amount = 3 * 200 + 1
                    timeAxis = np.linspace(minTime, maxTime,
                                           sample_points_amount, 
                                           endpoint=True)
                    resampled_data = np.interp(timeAxis, time, data)
                    dataDict[k] = resampled_data.tolist()
    # Check the data integrity 
    if is_sublist(["ia","ib", "ic"], dataDict.keys()):
        for idx, label in enumerate(labels):
            labels[idx] = np.append(np.zeros(200),
                                    np.ones(401) * int(label))
            labels[idx] = labels[idx].tolist()
        dataDict["label"] = labels
        dataDict["time"] = timeAxis.tolist() 
        dataDict["type"] = "simulation"
        # for easy to debug.
        dataDict["location"] = path
        try:
            col_sour.insert_one(dataDict)
        except:
            print(path)

def readFile(path:str, col:pymongo.collection.Collection):
    # our sample step.
    period = 0.02
    # this data can got from the header.
    # original sample rate
    with open(path, mode="r") as f:
        while True:
            line = f.readline()
            if "SampleRate" in line:
                break
    line = line.split(",")
    # read the frequency.
    source_freq = float(line[1].strip())
    period_points = int(source_freq * period)
    # get the labels
    path_splited = path.split("/")
    # start with . means is the hidden data
    if path_splited[-1][0] == ".":
        # print(path)
        return
    labels = []
    if path_splited[-1][0] == "0":
        labels.append("0")
    else:
        labels.append(path_splited[-1][0])
        # isdigit is a function for str.m
        if path_splited[-1][1].isdigit():
            labels.append(path_splited[-1][1])
    sour_data = np.loadtxt(path, skiprows=16, delimiter=",", usecols=range(1,4))
    # default sample point number of original data is (125000, 3)
    assert sour_data.shape == (125000, 3)
    # take oppsite number.
    sour_data = - sour_data
    # only using two periods data
    delimiter = int(sour_data.shape[0] / 2)
    # including the endpoint.
    try:
        data_slice = sour_data[delimiter-2*period_points:delimiter+1*period_points+1, :]
    except:
        print(period_points)
        return
    sample_points_amount = 3 * 200 + 1
    idx = np.linspace(0, 
                    data_slice.shape[0]-1,
                    sample_points_amount,
                    endpoint=True).astype(int)
    names = ["ia", "ib", "ic"]
    data_dict = {}
    for i, name in enumerate(names):
        data = data_slice[:,i]
        try:
            data_dict[name] = data[idx].tolist()
        except:
            raise
    for idx, label in enumerate(labels):
        labels[idx ] = np.append(np.zeros(2*200),
                                np.ones(201) * int(label))
        labels[idx] = labels[idx].tolist()
    data_dict["label"] = labels
    data_dict["type"] = "experiment"
    data_dict["location"] = path
    try:
        col.insert_one(data_dict)
    except:
        print(path)

def retrieve_files(path:str) -> Generator:
    path_gen = os.walk(path)
    for root, _, files in path_gen:
        for name in files:
           yield(os.path.join(root,name))

if __name__ == "__main__":
    scriptDir = os.path.dirname(os.path.realpath(__file__))
    basepath = os.path.dirname(os.path.dirname(scriptDir))
    basepath = os.path.dirname(basepath)
    dataDir = os.path.join(basepath, "data")
    exper_data_dir = os.path.join(dataDir, "new_data/并网逆变仿真数据/实验")
    simu_data_dir = os.path.join(dataDir, "new_data/并网逆变仿真数据")
    client = pymongo.MongoClient("mongodb://127.0.0.1:27017/")
    db  = client["Power_Fault"]
    col_sour = db["data_sour"]
    exper_file_gen = retrieve_files(exper_data_dir)
    for exper_file in exper_file_gen:
        if not fnmatch.fnmatch(exper_file, "*.csv"):
            continue
        try:
            readFile(exper_file, col_sour)
        except:
            break
    for i in range(1,5):
        split_data = os.path.join(dataDir, str(i))
        for tempName in os.listdir(split_data): 
                tempPath = os.path.join(split_data, tempName)
                # if item is a file skip it.
                if os.path.isfile(tempPath):
                    continue
                elif tempName == "1Normal":
                    readDirectory(tempPath, col_sour)
                else:
                    for sub_folder in os.listdir(tempPath):
                        sub_folder_path = os.path.join(tempPath, sub_folder)
                        readDirectory(sub_folder_path, col_sour)